function R = fmm_gibbs(gm, X, K, N, varargin)
% Performs Gibbs sampling of a finite mixture model
%
%   R = fmm_gibbs(gm, X, K, N, ...);
%       performs Gibbs sampling of a finite mixture model.
%
%       Inputs arguments:
%       - gm: the underlying generative model (such as gaussgm_gp)
%       - X:  the observed samples 
%       - K:  the number of component models
%       - N:  the number of samples to collect
%
%       The output R is a struct with the following fields:
%       - tag:  a string 'mm_samples' (indicating what the struct is for)
%       - alg:  a string 'fmm_gibbs' (indicating how it is derived)
%       - n:    the number of samples
%       - samples:  a struct array, of which each element 
%                   corresponds to a sample and has the following fields:
%                   - Id:       row vector, which is 1:K
%                   - Theta:    model parameters [param.dim x K]
%                   - Pi:       the prior distribution [1 x K]
%                   - W:        the posterior weights of each component
%
%       In addition, the caller can specify the following
%       options in name/value list:
%
%       - 'BurnIn':     the number of iterations to burn in 
%                       default = 200.
%
%       - 'Interval':   the number of iterations to take before collecting
%                       next sample. default = 20
%
%       - 'PriCount':   the count in Dirichlet prior for each component. 
%                       This can be either a vector of length K, or a
%                       scalar (if all components have the same count).
%                       default = 0.
%
%       Note that the option names are case-sensitive.
%

%   History
%   -------
%       - Created by Dahua Lin, on Nov 15, 2010
%

%% verify input arguments

if ~isobject(gm)
    error('fmm_gibbs:invalidarg', ...
        'gm should be an object to represent a generative model.');
end

ns = gm.check_observations(X);
if ns < 0
    error('fmm_gibbs:invalidarg', 'The observations X is not valid.');
end

K = check_pos_int('K', K);
N = check_pos_int('N', N);

opts = struct('BurnIn', 200, 'Interval', 20, 'PriCount', 0);
if ~isempty(varargin)
    opts = parlist(opts, varargin{:});
    
    opts.BurnIn = check_pos_int('BurnIn', opts.BurnIn);
    opts.Interval = check_pos_int('Interval', opts.Interval);
    
    pc = opts.PriCount;
    if ~(isnumeric(pc) && ...
            (isscalar(pc) || (isvector(pc) && numel(pc) == K) ) )
        error('fmm_gibbs:invalidarg', ...
            'PriCount should be a positive real value.');
    end
    if size(pc, 2) > 1; pc = pc.'; end  
    if ~isa(pc, 'double'); pc = double(pc); end
    opts.PriCount = pc;
else
    pc = opts.PriCount;
end


%% main

% initialize

if isscalar(opts.PriCount)
    p0 = constmat(K, 1, 1/K);
else
    p0 = (1 + pc) / (sum(pc) + K);
end
z = ddsample(p0, ns).';

% burn in

for t = 1 : opts.BurnIn
    [theta, pri, z] = fmm_move(gm, X, K, ns, pc, z);
end

% collect

samples = cell(N, 1);
samples{1} = make_sample(K, theta, pri, z, pc);

for i = 2 : N
    
    for t = 1 : opts.Interval
        [theta, pri, z] = fmm_move(gm, X, K, ns, pc, z);
    end
    
    samples{i} = make_sample(K, theta, pri, z, pc);  
end

% output

R.tag = 'mm_samples';
R.alg = 'fmm_gibbs';
R.n = N;
R.samples = vertcat(samples{:});


%% Core sampling step

function [theta, pri, z] = fmm_move(gm, X, K, ns, priCount, z)
% perform one MCMC iteration

spri = gm.npri == 1;

% theta ~ theta | z
theta = zeros(gm.pdim, K);
grp = intgroup(K, z);
for k = 1 : K
    sX = gm.select_observations(X, grp{k});
    if spri
        theta(:,k) = gm.pos_sample(sX, 1);
    else
        theta(:,k) = gm.pos_sample(sX, 1, 1, k);
    end
end

% pri ~ pri | z

c = intcount(K, z);
alpha = (1 + priCount) + c(:);
gv = randg(alpha);
pri = gv / sum(gv);

% z ~ z | theta, pri

LLik = gm.loglik(theta, X);
lpri = log(pri);
z = fmm_draw_labels(lpri, LLik, rand(1, ns));



%% Auxiliary functions

function s = make_sample(K, theta, pri, z, pc)

c = intcount(K, z);
s = struct( ...
    'Id', 1:K, ...
    'Theta', theta, ...
    'Pi', pri, ...
    'W', c + pc);


function v = check_pos_int(name, v)

if ~(isnumeric(v) && isscalar(v) && v > 0 && v == fix(v))
    error('fmm_gibbs:invalidarg', ...
        '%s should be a positive integer scalar.', name);
end
v = double(v);

